# This script reads in the time-dependent spin current correlation 
# functions and performs Fourier Transforms to obtain the AC spin 
# conductivity. It also performs checks for the sum rules.
#
# Author: Ehsan Khatami (San Jose State University, 2018)

import numpy as np
import matplotlib.pylab as plt

# Take the name of the file containing the current correlators, 
# e.g., "Wynn02_spin_current_01_site.dat"
file_curr = input('Enter file name for correlations: ')

# Take the name of the file containing the other properties,
# e.g., "Bare_properties_05_site.dat"
file_prop = input('Enter file name for other properties: ')

# Reading in the current correlation data
data_curr = np.loadtxt(file_curr,float)
times = data_curr[:100,1]
dt = times[1] - times[0]

# Read in the other properties
data_prop = np.loadtxt(file_prop,float)
T_values = data_prop[:,0]

# Taking the temperature
Tin = float(input('Enter the temperature value: '))

# Fine the temperature index
Tidx = np.argmin(abs(T_values-Tin))
T = T_values[Tidx]

# Taking the cut-off time for the F.T.
t_cut = float(input('Enter the time cut-off: '))

# Fine the time index for the cut-off
tidx = np.argmin(abs(times-t_cut))

print('Closest temperature used: ', T)
print('Closest cut-off time used: ', times[tidx])

# Getting the Re and Im parts of the current correlators at the 
# desired temperature
Re_Curr_Corr = data_curr[(Tidx-1)*100:Tidx*100,2]
Im_Curr_Corr = data_curr[(Tidx-1)*100:Tidx*100,3]

# Plotting the JJ correlation function vs time
plt.subplot(211)
plt.plot(times,Re_Curr_Corr, label='Re part')
plt.plot(times,Im_Curr_Corr, label='Im part')
plt.plot([t_cut,t_cut],[-0.1,0.1], '--k', label='$t_{cut-off}$')
plt.xlabel('Time [1/hopping]')
plt.ylabel('<$J_s(t)J_s(0)$> [hopping$^2$]')
plt.legend()
plt.grid(True)

# Getting the other properties at the desired temperature
entropy    = data_prop[Tidx-1,1]
kinetic    = data_prop[Tidx-1,2]
susc       = data_prop[Tidx-1,3]
cond_beta2 = data_prop[Tidx-1,4]

print('Entropy: ', entropy)
print('Spin susceptibility: ', susc)
print('Kinetic energy: ', kinetic)
print('Estimate of DC spin conductivity from imag-time <JJ>: ', cond_beta2)

# Performing the Fourier Transforms

# Number of frequency values
nws = 200
w = np.linspace(0.01,30,nws)
dw = w[1] - w[0]

# For the two ways of calculating the conductivity
FT1 = np.zeros(nws,dtype='complex')
FT2 = np.zeros(nws,dtype='complex')

for i in range(tidx+1):

	# Using the Simpsons rule
        if i == 0 or i == tidx:
           factor = 1./3
        elif i%2 == 1:
           factor = 4./3
        else:
           factor = 2./3

        FT1[:] += dt * factor * np.exp(1j*w[:]*times[i]) * (Re_Curr_Corr[i]+1j*Im_Curr_Corr[i])
        FT2[:] += dt * factor * np.exp(1j*w[:]*times[i]) * Im_Curr_Corr[i]

# Real parts of the AC spin conductivity calculated using the two different eqs.
Rsigma = FT1.real*(1-np.exp(-w/T)) / w
Isigma = -2 * FT2.imag / w

# Plotting the AC spin conductivity
plt.subplot(212)
plt.plot(w,Rsigma, label='Eq. (S17)')
plt.plot(w,Isigma, label='Eq. (S18)')
plt.xlabel('$\omega$ [hopping]')
plt.ylabel('Re $\sigma_s(\omega)$')
plt.legend()
plt.grid(True)
plt.tight_layout()

# Checking the sum rules
integRsigma = 0.0
integIsigma = 0.0
integRsigma_alt = 0.0
integIsigma_alt = 0.0

# Note that there is an extra factor of 2 in "alt" integrals
# due to considering only positive omega's
for i in range(nws):

     if i == 0 or i == nws-1:
        factor = 1./3
     elif i%2 == 1:
        factor = 4./3
     else:
        factor = 2./3

     integRsigma += dw * factor * Rsigma[i]
     integIsigma += dw * factor * Isigma[i]
     integRsigma_alt += dw * factor * w[i] * Rsigma[i] / np.sinh(w[i]/T/2.) / np.pi
     integIsigma_alt += dw * factor * w[i] * Isigma[i] / np.sinh(w[i]/T/2.) / np.pi

# Comparing integral to the x-component of KE (KE_x=KE/2)
print('Integral of Re sigma_s:', integRsigma, '-pi*<KE_x>/8:', -np.pi*kinetic/16)
print('Integral of Re sigma_s (alt):', integIsigma, '-pi*<KE_x>/8:', -np.pi*kinetic/16)

print('Lambda_s(beta/2)/pi/T^2$ (from integral):', \
      integRsigma_alt/T**2/np.pi, 'Direct:', cond_beta2)
print('Lambda_s(beta/2)/pi/T^2$ (from integral (alt)):', \
      integIsigma_alt/T**2/np.pi, 'Direct:', cond_beta2)

plt.show()
